use rig::client::{EmbeddingsClient, ProviderClient};
use rig::embeddings::EmbeddingsBuilder;
use rig::one_or_many::OneOrMany;
use rig::vector_store::in_memory_store::InMemoryVectorStore;
use rig::vector_store::VectorStoreIndex;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use tokio::sync::RwLock;

use crate::agent_types::AgentCategory;
use crate::agents::AgentManager;

/// Embedding document structure
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct EmbeddingDocument {
    /// Document ID
    pub id: String,
    /// Document content
    pub content: String,
    /// Document metadata
    pub metadata: HashMap<String, String>,
    /// Embedding vector (optional, will be computed if not provided)
    pub vector: Option<Vec<f32>>,
}

impl Eq for EmbeddingDocument {}

impl rig::Embed for EmbeddingDocument {
    fn embed(
        &self,
        _embedder: &mut rig::embeddings::TextEmbedder,
    ) -> Result<(), rig::embeddings::EmbedError> {
        Ok(())
    }
}

/// Embedding search result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingSearchResult {
    /// Document ID
    pub document_id: String,
    /// Document content
    pub content: String,
    /// Document metadata
    pub metadata: HashMap<String, String>,
    /// Similarity score
    pub score: f32,
}

/// Embedding manager for handling document embeddings and similarity search
pub struct EmbeddingManager {
    /// Vector store for embeddings
    vector_store: RwLock<InMemoryVectorStore<EmbeddingDocument>>,
    /// Cache for embedding models
    // We'll use a different approach since EmbeddingModel requires Sized
    _embedding_models: (),
}

fn create_empty_document() -> EmbeddingDocument {
    EmbeddingDocument {
        id: "empty".to_string(),
        content: "".to_string(),
        metadata: HashMap::new(),
        vector: None,
    }
}

impl std::fmt::Debug for EmbeddingManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EmbeddingManager")
            .field("vector_store", &"<RwLock<InMemoryVectorStore>>")
            .finish()
    }
}

impl EmbeddingManager {
    /// Create a new embedding manager
    pub fn new() -> Self {
        Self {
            vector_store: RwLock::new(InMemoryVectorStore::from_documents(vec![])),
            _embedding_models: (),
        }
    }

    /// Create an embedding model for a provider (simplified approach)
    fn create_embedding_model(
        &self,
        provider: &str,
        model: &str,
    ) -> anyhow::Result<rig::providers::openai::EmbeddingModel> {
        match provider {
            "openai" => {
                let client = rig::providers::openai::Client::from_env();
                Ok(client.embedding_model(model))
            }
            _ => Err(anyhow::anyhow!(
                "Unsupported provider for embeddings: {}",
                provider
            )),
        }
    }

    /// Add documents to the embedding store
    pub async fn add_documents(
        &self,
        provider: &str,
        model: &str,
        documents: Vec<EmbeddingDocument>,
    ) -> anyhow::Result<()> {
        let embedding_model = self.create_embedding_model(provider, model)?;

        let mut embeddings_builder = EmbeddingsBuilder::new(embedding_model);

        for doc in &documents {
            embeddings_builder = embeddings_builder.document(doc)?;
        }

        let embeddings = embeddings_builder.build().await?;

        // Create a new vector store with the documents and embeddings
        let mut vector_store = self.vector_store.write().await;
        let documents_for_store: Vec<_> = documents
            .clone()
            .into_iter()
            .zip(embeddings.into_iter())
            .map(|(doc, embedding)| {
                // Extract the embedding vector from OneOrMany
                let embedding_vec = embedding.1.first().vec;
                let rig_embedding = rig::embeddings::Embedding {
                    vec: embedding_vec,
                    document: doc.content.clone(),
                };
                (doc.id.clone(), doc.clone(), OneOrMany::one(rig_embedding))
            })
            .collect();

        *vector_store = InMemoryVectorStore::from_documents_with_ids(documents_for_store);

        Ok(())
    }

    /// Search for similar documents
    pub async fn search_similar(
        &self,
        provider: &str,
        model: &str,
        query: &str,
        top_k: usize,
    ) -> anyhow::Result<Vec<EmbeddingSearchResult>> {
        let embedding_model = self.create_embedding_model(provider, model)?;
        let vector_store = self.vector_store.read().await;
        let index = vector_store.clone().index(embedding_model);

        // Perform similarity search
        let search_request = rig::vector_store::VectorSearchRequest::builder()
            .query(query.to_string())
            .build()
            .map_err(|e| anyhow::anyhow!("Failed to build search request: {}", e))?;
        let results = index.top_n::<EmbeddingDocument>(search_request).await?;

        // Convert results to our format
        let search_results: Vec<EmbeddingSearchResult> = results
            .into_iter()
            .map(|(id, score, doc)| EmbeddingSearchResult {
                document_id: id.to_string(),
                content: doc.content,
                metadata: doc.metadata,
                score: score.parse().unwrap_or(0.0),
            })
            .collect();

        Ok(search_results)
    }

    /// Create embeddings for a single text
    pub async fn create_embedding(
        &self,
        provider: &str,
        model: &str,
        text: &str,
    ) -> anyhow::Result<Vec<f32>> {
        let embedding_model = self.create_embedding_model(provider, model)?;

        let temp_doc = EmbeddingDocument {
            id: "temp".to_string(),
            content: text.to_string(),
            metadata: HashMap::new(),
            vector: None,
        };
        let embeddings = EmbeddingsBuilder::new(embedding_model)
            .document(&temp_doc)?
            .build()
            .await?;

        // Extract the actual embedding vector from OneOrMany
        match embeddings.first() {
            Some((_, embedding)) => Ok(embedding.first().vec.iter().map(|&x| x as f32).collect()),
            None => Err(anyhow::anyhow!(
                "Failed to create embedding: no embedding returned"
            )),
        }
    }

    /// Get document by ID
    pub async fn get_document(
        &self,
        _document_id: &str,
    ) -> anyhow::Result<Option<EmbeddingDocument>> {
        let _vector_store = self.vector_store.read().await;
        // Note: This is a simplified implementation. In a real scenario,
        // you'd need to store the full document content separately or
        // extend the vector store to support document retrieval.
        Ok(None)
    }

    /// Delete document by ID
    pub async fn delete_document(&self, _document_id: &str) -> anyhow::Result<bool> {
        let _vector_store = self.vector_store.write().await;
        // Note: This is a simplified implementation. You'd need to implement
        // proper document deletion in the vector store.
        Ok(false)
    }

    /// List all document IDs
    pub async fn list_documents(&self) -> anyhow::Result<Vec<String>> {
        let _vector_store = self.vector_store.read().await;
        // Note: This is a simplified implementation. You'd need to implement
        // proper document listing in the vector store.
        Ok(vec![])
    }
}

/// Create an embedding agent for knowledge base operations
pub async fn create_embedding_agent(
    agent_manager: &AgentManager,
    embedding_manager: &EmbeddingManager,
    name: &str,
    provider: &str,
    model: &str,
    preamble: &str,
    documents: Vec<EmbeddingDocument>,
) -> anyhow::Result<i64> {
    // Create the agent configuration
    let agent_id = agent_manager
        .create_agent(
            name,
            provider,
            model,
            preamble,
            None,
            None,
            AgentCategory::Embed,
            None,
        )
        .await?;

    // Add documents to the embedding store
    if !documents.is_empty() {
        embedding_manager
            .add_documents(provider, model, documents)
            .await?;
    }

    Ok(agent_id)
}

/// Perform RAG (Retrieval-Augmented Generation) with an embedding agent
pub async fn rag_query(
    agent_manager: &AgentManager,
    embedding_manager: &EmbeddingManager,
    agent_id: i64,
    query: &str,
    context_docs: usize,
) -> anyhow::Result<String> {
    // Get the agent configuration
    let agent = agent_manager
        .get_agent(agent_id)
        .await?
        .ok_or_else(|| anyhow::anyhow!("Agent not found"))?;

    // Search for relevant documents
    let search_results = embedding_manager
        .search_similar(
            &agent.base_config.provider,
            &agent.base_config.model,
            query,
            context_docs,
        )
        .await?;

    // Build context from search results
    let context: String = search_results
        .iter()
        .map(|result| {
            format!(
                "Document: {}\nContent: {}\n",
                result.document_id, result.content
            )
        })
        .collect::<Vec<_>>()
        .join("\n---\n");

    // Create enhanced prompt with context
    let enhanced_prompt = if context.is_empty() {
        query.to_string()
    } else {
        format!(
            "Context:\n{}\n\nQuestion: {}\n\nPlease answer the question based on the provided context.",
            context, query
        )
    };

    // Use the completion module to get the response
    let temp_dir = tempfile::TempDir::new().unwrap();
    let tool_manager = crate::tools::ToolManager::new(temp_dir.path().join("tools.json")).unwrap();
    crate::completion::prompt(agent_manager, &tool_manager, agent_id, &enhanced_prompt).await
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    #[tokio::test]
    async fn test_embedding_manager_creation() {
        let manager = EmbeddingManager::new();
        assert!(manager.list_documents().await.is_ok());
    }

    #[tokio::test]
    async fn test_create_embedding_agent() {
        let temp_dir = TempDir::new().unwrap();
        let agents_file = temp_dir.path().join("agents.json");
        let agent_manager = AgentManager::new(&agents_file).unwrap();
        let embedding_manager = EmbeddingManager::new();

        let documents = vec![EmbeddingDocument {
            id: "doc1".to_string(),
            content: "This is a test document about embeddings.".to_string(),
            metadata: HashMap::new(),
            vector: None,
        }];

        let result = create_embedding_agent(
            &agent_manager,
            &embedding_manager,
            "Test Embed Agent",
            "openai",
            "text-embedding-ada-002",
            "You are an embedding assistant.",
            documents,
        )
        .await;

        assert!(result.is_ok());
    }
}
